import sys
import os
from volcenginesdkarkruntime import Ark
from dotenv import load_dotenv
load_dotenv()

client = Ark()
# model_id = "ep-20240522112440-fmmfs" #doubao-pro-4k
model_id = "ep-20241224120132-x84bs"  # Doubao-pro-32k-241215


def call_doubao_api(prompt):
    completion = client.chat.completions.create(
        model=model_id,
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant."
            },
            {
                "role": "user",
                "content": prompt
            }
        ],
        stream=False,
        temperature=0
    )
    try:
        text_response = completion.choices[0].message.content.strip()
        return text_response

        # token_cost = completion.usage['total_tokens']
        # money_cost = (token_cost/1000)*0.001
        # print(f"消耗token:{token_cost}，耗费钱：{money_cost}")
        return text_response
    except Exception as e:
        return "Doubao call exception"
